import gradio as gr

from ai_translator_v2 import PDFTranslator
from ai_translator_v2.translator.translation_config import TranslationConfig
from ai_translator_v2.utils import ArgumentParser


def translate(text, style, source_language, target_language):
    args = ArgumentParser().parse_args()
    config = TranslationConfig()
    config.initialize(args)
    return PDFTranslator(config.model_name).translate_pdf(
        text,
        style,
        source_language=source_language,
        target_language=target_language)


s_text = gr.Text(label='源语言', value="English", placeholder="请输入源语言")
t_text = gr.Text(label='目标语言', value="Chinese", placeholder="请输入目标语言")
style_select = gr.Dropdown(label="风格选择", choices=['作家', '诗人', '文学家', '翻译家'], value='诗人')
interFace = gr.Interface(fn=translate, inputs=['file', style_select, s_text, t_text], outputs="file")
if __name__ == "__main__":
    interFace.launch()
